#include <cstdio>
#include <cstring>
#include "cuda.h"
#include <algorithm>
#include <sys/mman.h>
#include "common.h"
#include <cmath>
using namespace std;

#define module_name "fft"

const int BLOCK_SIZE = 256;

CUmodule cuModule;

int ceil_div(int a, int b) { return (a+b-1)/b; }

typedef pair<double,double> K;

void fft(CUdeviceptr arr, int P, bool rev) {
	const int BLOCK_NUM_X = ceil(sqrt(ceil_div(P, BLOCK_SIZE)));
	const int BLOCK_NUM_Y = ceil(sqrt(ceil_div(P, BLOCK_SIZE)));
	// permutate
	CUfunction permutate;
	cuCall(cuModuleGetFunction(&permutate, cuModule, "permutate"), "cannot acquire kernel permutate");
	void *args_perm[] = {&arr, &P};
	cuCall(cuLaunchKernel(permutate, BLOCK_NUM_X, BLOCK_NUM_Y, 1, BLOCK_SIZE, 1, 1, 0, 0, args_perm, 0), "cannot run kernel permutate");

	int n_step = 0, step;
	while((1<<n_step) < P) n_step++;
	CUfunction fft;
	cuCall(cuModuleGetFunction(&fft, cuModule, "fft"), "cannot acquire kernel fft");
	void *args[] = {&arr, &P, &step, &n_step, &rev};
	for(step=1;step<P;step*=2)
	{
		n_step--;
		cuCall(cuLaunchKernel(fft, BLOCK_NUM_X, BLOCK_NUM_Y, 1, BLOCK_SIZE, 1, 1, 0, 0, args, 0), "cannot run kernel fft");
	}
}

int multiply(unsigned char* Astr, int n, unsigned char* Bstr, int m, char* Cstr) {
	// load module blah blah
	cuCall(cuModuleLoad(&cuModule, "./" module_name ".ptx"), "cannot load module");

	// prepare the numbers; reverse; change the base
	CUfunction reverse_num;
	cuCall(cuModuleGetFunction(&reverse_num, cuModule, "reverse_num"), "cannot acquire kernel reverse_num");
	void* args_reverse_num[] = {&Astr, &n, &Bstr, &m};
	const int BLOCK_NUM_X_REV_NUM = ceil(sqrt(ceil_div(max(n, m), BLOCK_SIZE)));
	const int BLOCK_NUM_Y_REV_NUM = ceil(sqrt(ceil_div(max(n, m), BLOCK_SIZE)));
	cuCall(cuLaunchKernel(reverse_num, BLOCK_NUM_X_REV_NUM, BLOCK_NUM_Y_REV_NUM, 1, BLOCK_SIZE, 1, 1, 0, 0, args_reverse_num, 0), "cannot run kernel reverse_num");

	CUdeviceptr A, B, C;
	const int A_DIGS = ceil_div(n, DIGS);
	const int B_DIGS = ceil_div(m, DIGS);
	unsigned P; for(P=1;P<A_DIGS || P<B_DIGS;P*=2); P *= 2;
	const int BLOCK_NUM_X = ceil(sqrt(ceil_div(P, BLOCK_SIZE)));
	const int BLOCK_NUM_Y = ceil(sqrt(ceil_div(P, BLOCK_SIZE)));
	cuCall(cuMemAlloc(&A, P * sizeof(K)), "cuMemAlloc(A)");
	cuCall(cuMemAlloc(&B, P * sizeof(K)), "cuMemAlloc(B)");
	CUfunction init;
	cuCall(cuModuleGetFunction(&init, cuModule, "init"), "cannot acquire kernel init");
	void* args_init[] = {&A, &Astr, &n, &B, &Bstr, &m, &P };
	cuCall(cuLaunchKernel(init, BLOCK_NUM_X, BLOCK_NUM_Y, 1, BLOCK_SIZE, 1, 1, 0, 0, args_init, 0), "cannot run kernel init");

	// multiply
	fft(A, P, false);
	fft(B, P, false);
	CUfunction multiply;
	cuCall(cuModuleGetFunction(&multiply, cuModule, "multiply"), "cannot acquire kernel multiply");
	void *args_multiply[] = {&A, &B, &P};
	cuCall(cuLaunchKernel(multiply, BLOCK_NUM_X, BLOCK_NUM_Y, 1, BLOCK_SIZE, 1, 1, 0, 0, args_multiply, 0), "cannot run kernel multiply");
	fft(A, P, true);

	// cos tam
	cuCall(cuMemFree(B), "cuMemFree(B)");
	cuCall(cuMemAlloc(&C, P * sizeof(double)), "cuMemAlloc(C)");
	CUfunction result;
	cuCall(cuModuleGetFunction(&result, cuModule, "result"), "cannot acquire kernel result");
	void *args_result[] = {&A, &C, &P};
	cuCall(cuLaunchKernel(result, BLOCK_NUM_X, BLOCK_NUM_Y, 1, BLOCK_SIZE, 1, 1, 0, 0, args_result, 0), "cannot run kernel result");
	cuCtxSynchronize();

	// TODO to moze dac sie zrobic szybko na gpu
	void *C_host_v;
	cuMemAllocHost(&C_host_v, P * sizeof(double));
	cuMemcpyDtoH(C_host_v, C, P * sizeof(double));
	double *C_host = (double*)C_host_v;
	long long s = 0;
	void *Cstr_host_v;
	cuMemAllocHost(&Cstr_host_v, DIGS * P * sizeof(char));

	unsigned char *Cstr_host = (unsigned char*)Cstr_host_v;
	long long k = 0;
	for(int i=0;i<P;i++)
	{
		s += C_host[i];
		long long t = s % BASE;
		s /= BASE;
		for(int j=0;j<DIGS;j++)
		{
			unsigned char foo = t % BASE_BASE;
			if (foo) k = DIGS*i+j;
			Cstr_host[DIGS*i+j] = (foo);
			t /= BASE_BASE;
		}
	}
	k++;
	CUfunction reverse;
	cuCall(cuModuleGetFunction(&reverse, cuModule, "reverse"), "cannot acquire kernel reverse");
	const int BLOCK_NUM_X_REV = ceil(sqrt(ceil_div(n + m, BLOCK_SIZE)));
	const int BLOCK_NUM_Y_REV = ceil(sqrt(ceil_div(n + m, BLOCK_SIZE)));
	void *args_reverse[] = {&Cstr_host_v, &Cstr, &k};
	cuCall(cuLaunchKernel(reverse, BLOCK_NUM_X_REV, BLOCK_NUM_Y_REV, 1, BLOCK_SIZE, 1, 1, 0, 0, args_reverse, 0), "cannot run kernel reverse");
	return k;
}

